from abc import ABC, abstractmethod
from contextlib import asynccontextmanager
import asyncio
from venv import logger
import asyncpg
from pgvector.asyncpg import register_vector
from ..common import DatabaseConfig, PostgresConfigurationSettings, get_logger

logger=get_logger(__name__)

class Handler(ABC):
    """数据库处理器抽象基类"""
    
    def __init__(self, project_name: str, connection_manager):
        """
        初始化处理器
        
        Args:
            project_name: 项目名称
            connection_manager: 数据库连接管理器（SemaphoreConnectionPool实例）
        """
        self.project_name = project_name
        self.connection_manager = connection_manager
    
    def _get_table_name(self, table_type: str) -> str:
        """
        获取表名
        
        Args:
            table_type: 表类型
            
        Returns:
            str: 完整的表名
        """
        return f"{self.project_name}_{table_type}"
    
    @abstractmethod
    async def create_tables(self) -> None:
        """
        创建数据库表
        
        Raises:
            NotImplementedError: 子类必须实现此方法
        """
        raise NotImplementedError("子类必须实现create_tables方法")


class BaseDatabaseManger(ABC):
    def __init__(self, db_config: DatabaseConfig):
        self.provider=db_config.provider
        self.db_config = db_config
        self.driver=db_config.driver
        self.user=db_config.user
        self.host=db_config.host
        self.port=db_config.port
        self.password=db_config.password
        self.db_name=db_config.db_name
    
    @abstractmethod
    async def initialize():
        pass  

    @abstractmethod
    async def get_connection():
        pass
    
    @abstractmethod
    async def close():
        pass

class SemaphoreConnectionPool:
    """使用信号量控制的异步PostgreSQL连接池，支持向量数据库操作"""
    
    def __init__(self, connection_string: str,
                postgres_configuration_settings: PostgresConfigurationSettings):
        self.connection_string = connection_string
        self.postgres_configuration_settings = postgres_configuration_settings
        self.pool = None
        self.semaphore = None

    async def initialize(self):
        """初始化连接池并注册向量类型支持"""
        try:
            logger.info(
                f"Connecting with {int(self.postgres_configuration_settings.max_connections * 0.9)} connections to `asyncpg.create_pool`."
            )

            self.semaphore = asyncio.Semaphore(
                int(self.postgres_configuration_settings.max_connections * 0.9)
            )

            self.pool = await asyncpg.create_pool(
                self.connection_string,
                max_size=self.postgres_configuration_settings.max_connections,
                statement_cache_size=self.postgres_configuration_settings.statement_cache_size,
                init=self._init_connection  # 添加连接初始化函数
            )

            logger.info(
                "Successfully connected to Postgres database and created connection pool."
            )
        except Exception as e:
            raise ValueError(
                f"Error {e} occurred while attempting to connect to relational database."
            ) from e
    
    @asynccontextmanager
    async def get_connection(self):
        """
        Get database connection with semaphore concurrency control
        
        This ensures concurrent connections do not exceed 90% of configured maximum,
        preventing system overload. Error handling is added to ensure connection
        exceptions are properly logged and propagated
        """
        async with self.semaphore:
            try:
                async with self.pool.acquire() as conn:
                    yield conn
            except Exception as e:
                logger.error(f"Error acquiring or using database connection: {e}")
                raise
            
    async def _init_connection(self, conn):
        """初始化单个连接，注册向量类型支持"""
        await register_vector(conn)

    async def close(self):
        """关闭连接池"""
        if self.pool:
            await self.pool.close()